{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Change Detection (Classification) using TAO Visual ChangeNet-Classification Network\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n",
    "\n",
    "## What is a Visual ChangeNet Network?\n",
    "\n",
    "Visual ChangeNet is a state of the art transformer-based Change Detection model. Similar to the optical inspection notebook, Visual ChangeNet is also based on Siamese Network, which is a class of neural network architectures containing two or more identical subnetworks. The training algorithm works by updating the parameters across all the sub-networks in tandem. In TAO, Visual ChangeNet supports two images as input where the end goal is to either classify or segment the change between the \"golden or reference\" image and the \"test\" image. TAO supports the [FAN](https://arxiv.org/abs/2204.12451) backbone network for both Visual ChangeNet architectures. For more details about training FAN backbones, please refer to the classification pytorch notebook.\n",
    "\n",
    "TAO Toolkit versions 5.3 and later support some of the foundational models for change detection (classification and segmentation). NV-DINOv2 can now be used as the backbone for the Visual ChangeNet-Classification and Segmentation models.\n",
    "\n",
    "In TAO, two different types of Change Detection networks are supported: Visual ChangeNet-Segmentation and Visual ChangeNet-Classification intended for classification and segmentation of change between the two input images, respectively. Visual ChangeNet-Classification is specifically intended for change classification. In this notebook, we leverage the Visual ChangeNet-Classification model to demonstrate PCB component Inspection. \n",
    "\n",
    "### Sample inference from the Visual ChangeNet-Classification model\n",
    "\n",
    "| **Pass** | **Fail** |\n",
    "| :--: | :--: |\n",
    "|<img align=\"center\" title=\"no defects\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/optical_inspection/pass.png?raw=true\" width=\"400\" > no defects | <img align=\"center\" title=\"missing component\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/optical_inspection/defect.png?raw=true\" width=\"400\" >  missing component |\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Train a Visual ChangeNet-Classification Model for Optical Inspection on a PCB Dataset.\n",
    "* Explore options to combine images of multiple LED light sources within Visual ChangeNet-Classification Network \n",
    "* Evaluate the trained model.\n",
    "* Run Inference on the trained model.\n",
    "* Export the trained model to a .onnx file (encrypted ONNX model) for deployment to DeepStream or TensorRT.\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `visual_changenet` model, \n",
    "which you may deploy with this [end-to-end sample](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps) with Triton.\n",
    "\n",
    "Visual ChangeNet-Classification TAO Features:\n",
    "\n",
    "* Combining multiple LED intensities, camera angles or different sensory inputs for image-pair comparison within Siamese Network\n",
    "    * Multiple Inputs (1....N)\n",
    "* Mode of concat\n",
    "    * Linear concat (1 X N)\n",
    "    * Grid concat (N X M)\n",
    "* Backbone networks supported for Visual ChangeNet-Classification\n",
    "    * FAN (FAN-T-Hybrid, FAN-S-Hybrid, FAN-B-Hybrid, FAN-L-Hybrid)\n",
    "    * NVDinoV2\n",
    "* Backbone Initialization\n",
    "    * Load Pre-trained Weights\n",
    "    * Train from Scratch\n",
    "\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of Siamese Network using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset](#head-2)\n",
    "3. [Download the pretrained model from NGC](#head-3)\n",
    "4. [Provide training specification](#head-4)\n",
    "5. [Run TAO training](#head-5)\n",
    "6. [Evaluate trained models](#head-6)\n",
    "7. [Inferences](#head-7)\n",
    "8. [Deploy](#head-8)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=FIXME\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\", \"changenet-classify\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"changenet-classify\", \"results\")\n",
    "os.environ[\"HOST_MODEL_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"changenet-classify\", \"model\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=/path/to/local/tao-experiments/changenet\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR\n",
    "! mkdir -p $HOST_MODEL_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tlt_configs = {\n",
    "   \"Mounts\":[\n",
    "       # Mapping the data directory\n",
    "       {\n",
    "           \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "           \"destination\": \"/workspace/tao-experiments\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_MODEL_DIR\"],\n",
    "           \"destination\": \"/model\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       }\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         }\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tlt_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.7.0. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7.0 < 3.8.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 525+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python >=3.7.0. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare a sample dataset <a class=\"anchor\" id=\"head-2\"></a>\n",
    "\n",
    "TAO Visual ChangeNet-Classification uses a custom dataset format, similar to the one used by the Optical Inspection model. The sections below walk through this format, specifically how the dataset should be structured and the files required."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visual ChangeNet-Classification data format\n",
    "\n",
    "Visual ChangeNet-Classification expects directories of images and csv files in the dataset root directory. The image directory consists of golden images (non-defective reference images) and test images to be\n",
    "compared with golden samples for PCB defect classification.  \n",
    "\n",
    "\n",
    "    |--dataset_root:\n",
    "       |--images\n",
    "            |--input:\n",
    "               |--C1.jpg\n",
    "               |--C2.jpg\n",
    "            |--golden:\n",
    "               |--C1.jpg\n",
    "               |--C2.jpg\n",
    "            |--input1:\n",
    "               |--C1.jpg\n",
    "               |--C3.jpg\n",
    "            |--golden1:\n",
    "               |--C1.jpg\n",
    "               |--C3.jpg\n",
    "       |--labels\n",
    "           |-- train.csv\n",
    "           |-- validation.csv\n",
    "\n",
    "Here's a description of the structure:\n",
    "\n",
    "* The images directory contains the following:\n",
    "    * input: Directory containing input images to be compared with golden images\n",
    "    * golden: Directory containing golden reference images\n",
    "* The labels directory contains the csv files for pair-wise images input to the Visual ChangeNet-Classification model with corresponding class labels. Details of this file are included in the `Label Files` section below. \n",
    "\n",
    "#### Label Files\n",
    "\n",
    "A Visual ChangeNet-Classification label file is a simple csv file containing the following fields.\n",
    "\n",
    "| **input_path** | **golden_path** | **label** | **object_name** |\n",
    "| :-- | :-- | :--- | :--- |\n",
    "| `/path/to/input/image/directory` | `/path/to/golden/image/directory` | The class to which the object belongs | `/component/name` |\n",
    "\n",
    "Here's a description of the structure:\n",
    "\n",
    "* ``input_path``: The path to the directory containing input compare image\n",
    "* ``golden_path``: The path to the directory containing corresponding golden reference image\n",
    "* ``label``: The labels for the pair-wise images (Use `PASS` for non-defective components, and any other specific defect type label for defective components).\n",
    "* ``object_name``: The name of the component to be compared. The object name is the same for input and golden images and represents the image name without the file extension.\n",
    "    * For each :code:`object_name`, TAO supports combining multiple LED intensities, camera angles or different sensory inputs for each of the input and golden images to be\n",
    "      compared within the Visual ChangeNet-Classification model. For more details, refer to the [Data Annotation Format](https://docs.nvidia.com/tao/tao-toolkit/text/data_annotation_format.html#id3) section.\n",
    "\n",
    "Here is a sample label file corresponding to the sample directory structure as described in the [Optical Inspection Format](https://docs.nvidia.com/tao/tao-toolkit/text/data_annotation_format.html#optical-inspection-format) section.\n",
    "\n",
    "| **input_path** | **golden_path** | **label** | **object_name** |\n",
    "| :-- | :-- | :--- | :--- |\n",
    "| /dataset_root/images/input/  | /dataset_root/images/golden/  | PASS      | C1  |\n",
    "| /dataset_root/images/input/  | /dataset_root/images/golden/  | PASS      | C2  |\n",
    "| /dataset_root/images/input1/ | /dataset_root/images/golden1/ | MISSING   | C1  |\n",
    "| /dataset_root/images/input1/ | /dataset_root/images/golden1/ | PASS      | C3  |\n",
    "\n",
    "``Note``: In the label file, ensure that non-defective samples are consistently labeled as PASS, \n",
    "while defective samples can be assigned any specific defect type label. \n",
    "The model is designed to treat all defects collectively and train for binary defect classification.\n",
    "\n",
    "#### Input Mapping\n",
    "\n",
    "TAO supports combining several lighting conditions (1...N) for each image-golden pair to be combined for comparison within the Visual ChangeNet-Classification Network. The following modes are supported: \n",
    "\n",
    "* Linear concat (1 x N)\n",
    "* Grid concat (M x N)\n",
    "\n",
    "Each `object_name` is appended with name of the lighting condition as specified in the experiment spec under `input_map`. Here is an example of dataset experiment spec changes for combining 4 input lighting conditions as a 1x4 linear grid:\n",
    "\n",
    "```yaml\n",
    "    dataset: \n",
    "        num_input: 4\n",
    "        concat_type: linear\n",
    "        input_map:\n",
    "            LowAngleLight: 0\n",
    "            SolderLight: 1\n",
    "            UniformLight: 2\n",
    "            WhiteLight: 3\n",
    "```\n",
    "\n",
    "The dataset also supports single lighting condition per input-golden image pair as input to the Visual ChangeNet-Classification Network. Here is an example of dataset experiment spec changes for single lighting condition where `object_name` represents the image name:\n",
    "\n",
    "```yaml\n",
    "    dataset: \n",
    "        num_inputs: 1\n",
    "        input_map: null\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Verify\n",
    "!ls -l $HOST_DATA_DIR/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Download the pretrained model from NGC <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "Download the pretrained model from NGC. We will use NGC CLI to get the data and model. For more details, go to https://ngc.nvidia.com and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "import os\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $HOST_RESULTS_DIR/ngccli\n",
    "\n",
    "# # Remove any previously existing CLI installations\n",
    "!rm -rf $HOST_RESULTS_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $HOST_RESULTS_DIR/ngccli\n",
    "!unzip -u \"$HOST_RESULTS_DIR/ngccli/$CLI\" -d $HOST_RESULTS_DIR/ngccli/\n",
    "!rm $HOST_RESULTS_DIR/ngccli/*.zip\n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"HOST_RESULTS_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/visual_changenet_classification:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/pretrained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model download-version \"nvidia/tao/visual_changenet_classification:visual_changenet_nvpcb_trainable_v1.0\" --dest $HOST_RESULTS_DIR/pretrained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $HOST_RESULTS_DIR/pretrained/visual_changenet_classification_vvisual_changenet_nvpcb_trainable_v1.0/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Provide training specification <a class=\"anchor\" id=\"head-4\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    "\n",
    "* task: classify\n",
    "* train:\n",
    "  * resume_training_checkpoint_path: null\n",
    "  * classify:\n",
    "    * loss: \"ce\"\n",
    "    * cls_weight: [1.0, 10.0]\n",
    "  * num_epochs: 30\n",
    "  * num_nodes: 1\n",
    "  * validation_interval: 1\n",
    "  * checkpoint_interval: 1\n",
    "  * optim:\n",
    "    * lr: 0.00005\n",
    "    * optim: \"adamw\"\n",
    "    * policy: \"linear\" \n",
    "    * momentum: 0.9\n",
    "    * weight_decay: 0.01\n",
    "    * betas: [0.9, 0.999]\n",
    "  * results_dir: \"${results_dir}/train\"\n",
    "  * tensorboard:\n",
    "    * enabled: True\n",
    "* results_dir: \"???\"\n",
    "* model:\n",
    "  * backbone:\n",
    "    * type: \"fan_small_12_p4_hybrid\"\n",
    "    * pretrained_backbone_path: /path/to/pretrained/backbone.pth\n",
    "    * freeze_backbone: False\n",
    "  * decode_head:\n",
    "    * feature_strides: [4, 8, 16, 16]\n",
    "  * classify:\n",
    "    * train_margin_euclid: 2.0\n",
    "    * eval_margin: 0.005\n",
    "    * embedding_vectors: 5\n",
    "    * embed_dec: 30\n",
    "    * difference_module: 'learnable'\n",
    "    * learnable_difference_modules: 4\n",
    "* dataset:\n",
    "  * classify:\n",
    "    * train_dataset:\n",
    "      * csv_path: /data/dataset_convert/train_combined.csv\n",
    "      * images_dir: /data/images/\n",
    "    * validation_dataset:\n",
    "      * csv_path: /data/dataset_convert/valid_combined.csv\n",
    "      * images_dir: /data/images/\n",
    "    * test_dataset:\n",
    "      * csv_path: /data/dataset_convert/valid_combined.csv\n",
    "      * images_dir: /data/images/\n",
    "    * infer_dataset:\n",
    "      * csv_path: /data/dataset_convert/valid_combined.csv\n",
    "      * images_dir: /data/images/\n",
    "    * image_ext: .jpg\n",
    "    * batch_size: 4\n",
    "    * workers: 1\n",
    "    * fpratio_sampling: 0.1\n",
    "    * num_input: 4\n",
    "    * input_map:\n",
    "      * LowAngleLight: 0\n",
    "      * SolderLight: 1\n",
    "      * UniformLight: 2\n",
    "      * WhiteLight: 3\n",
    "    * concat_type: linear\n",
    "    * grid_map:\n",
    "      * x: 2\n",
    "      * y: 2\n",
    "    * image_width: 128\n",
    "    * image_height: 128\n",
    "    * augmentation_config:\n",
    "      * rgb_input_mean: [0.485, 0.456, 0.406]\n",
    "      * rgb_input_std: [0.229, 0.224, 0.225]\n",
    "      * random_flip:\n",
    "        * vflip_probability: 0.5\n",
    "        * hflip_probability: 0.5\n",
    "        * enable: True\n",
    "      * random_rotate:\n",
    "        * rotate_probability: 0.5\n",
    "        * angle_list: [90, 180, 270]\n",
    "        * enable: True\n",
    "      * random_color:\n",
    "        * brightness: 0.3\n",
    "        * contrast: 0.3\n",
    "        * saturation: 0.3\n",
    "        * hue: 0.3\n",
    "        * enable: True\n",
    "      * with_scale_random_crop:\n",
    "        * enable: True\n",
    "      * with_random_crop: True\n",
    "      * with_random_blur: True\n",
    "      * augment: False\n",
    " \n",
    "Please refer to the TAO documentation about Visual ChangeNet-Classification to get all the parameters that are configurable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/experiment_classify.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Run TAO training <a class=\"anchor\" id=\"head-5\"></a>\n",
    "* Provide the sample spec file and the output directory location for models.\n",
    "* WARNING: Training will take several hours or one day to complete.\n",
    "* Training time on a single GPU for an internal dataset with (~40,000 training images) (the hours may vary depending on the data location, network speed, and etc).\n",
    "\n",
    "<table>\n",
    "  <tr>\n",
    "    <th>GPU Type</th>\n",
    "    <th>Time (minutes)</th>\n",
    "  </tr>\n",
    "  <tr>\n",
    "    <td>1 x A6000 48GB</td>\n",
    "    <td>15</td>\n",
    "  </tr>\n",
    "</table>\n",
    "\n",
    "* If you're looking to speed through the notebook, you can adjust the `train.num_epochs` parameter for a quicker run. It's crucial to understand that this adjustment is intended for demonstration and educational purposes only and won't produce state-of-the-art (SOTA) results.\n",
    "* To replicate baseline numbers accurately, the model should be trained for 50 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env MODEL_DIR = /model\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results\n",
    "%env NUM_EPOCHS=5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.1 Train Visual ChangeNet-Classification model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will train a Visual ChangeNet-Classification model from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Train model\")\n",
    "!tao model visual_changenet train \\\n",
    "                  -e $SPECS_DIR/experiment_classify.yaml \\\n",
    "                  train.pretrained_model_path=$RESULTS_DIR/pretrained/visual_changenet_classification_vvisual_changenet_nvpcb_trainable_v1.0/changenet_classifier.pth\n",
    "                #   train.num_epochs=$NUM_EPOCHS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Training command for multi-gpu training. We can define the number of gpus using the `train.num_gpus` parameter.\n",
    "## The following command will trigger multi-gpu training on 2 gpus.\n",
    "# print(\"Train model\")\n",
    "# !tao model visual_changenet train \\\n",
    "#                   -e $SPECS_DIR/experiment_classify.yaml \\\n",
    "#                   train.num_gpus 2\n",
    "#                   train.num_epochs=$NUM_EPOCHS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Model checkpoints:')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"train/changenet_model_classify_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/train/changenet_classify.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train/changenet_classify.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Evaluate trained models <a class=\"anchor\" id=\"head-6\"></a>\n",
    "Evaluate trained model.\n",
    "\n",
    "For Visual ChangeNet-Classification model evaluation, we use the following metrics:\n",
    "\n",
    "* Defect Accuracy: Defect detection accurary.\n",
    "* False Positive Rate (FPR): False Alarm Rate for for non-defective samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model visual_changenet evaluate \\\n",
    "                   -e $SPECS_DIR/experiment_classify.yaml \\\n",
    "                    evaluate.checkpoint=$RESULTS_DIR/train/changenet_classify.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Inferences <a class=\"anchor\" id=\"head-7\"></a>\n",
    "In this section, we run the Visual ChangeNet-Classification inference tool to generate inferences with the trained models and save the results under `$RESULTS_DIR`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model visual_changenet inference \\\n",
    "                   -e $SPECS_DIR/experiment_classify.yaml \\\n",
    "                    inference.checkpoint=$RESULTS_DIR/train/changenet_classify.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Deploy <a class=\"anchor\" id=\"head-8\"></a>\n",
    "Export the model to encrypted ONNX model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model visual_changenet export \\\n",
    "                    -e $SPECS_DIR/experiment_classify.yaml \\\n",
    "                        export.checkpoint=$RESULTS_DIR/train/changenet_classify.pth \\\n",
    "                        export.onnx_file=$RESULTS_DIR/export/changenet-classify.onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # # Uncomment this cell to export an onnx file with dynamic batching enabled for\n",
    "# # # integration with trtexec and deepstream.\n",
    "# # # Export the model with export.batch_size=-1 for dynamic batching.\n",
    "# ! tao model visual_changenet export \\\n",
    "#                     -e $SPECS_DIR/experiment_classify.yaml \\\n",
    "#                     export.checkpoint=$RESULTS_DIR/train/changenet_classify.pth \\\n",
    "#                     export.batch_size=-1 \\\n",
    "#                     export.results_dir=$RESULTS_DIR/export_dynamic \\\n",
    "#                     export.onnx_file=$RESULTS_DIR/export_dynamic/changenet-classify.onnx\n",
    "\n",
    "# # Profiling the exported model via trtexec.\n",
    "# ! tao deploy trtexec --onnx=$RESULTS_DIR/export_dynamic/changenet-classify.onnx \\\n",
    "#                      --minShapes=input_1:1x3x512x128,input_2:1x3x512x128 \\\n",
    "#                      --optShapes=input_1:8x3x512x128,input_2:8x3x512x128 \\\n",
    "#                      --maxShapes=input_1:16x3x512x128,input_2:16x3x512x128 \\\n",
    "#                      --fp16 \\\n",
    "#                      --saveEngine=$RESULTS_DIR/export_dynamic/changenet-classify-fp16.engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy visual_changenet gen_trt_engine \\\n",
    "            -e $SPECS_DIR/experiment_classify.yaml \\\n",
    "                gen_trt_engine.onnx_file=$RESULTS_DIR/export/changenet-classify.onnx \\\n",
    "                gen_trt_engine.trt_engine=$RESULTS_DIR/gen_trt_engine/changenet-classify.trt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy visual_changenet inference \\\n",
    "            -e $SPECS_DIR/experiment_classify.yaml \\\n",
    "                inference.trt_engine=$RESULTS_DIR/gen_trt_engine/changenet-classify.trt \\\n",
    "                inference.results_dir=$RESULTS_DIR/trt_inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy visual_changenet evaluate \\\n",
    "            -e $SPECS_DIR/experiment_classify.yaml \\\n",
    "                evaluate.trt_engine=$RESULTS_DIR/gen_trt_engine/changenet-classify.trt \\\n",
    "                evaluate.results_dir=$RESULTS_DIR/trt_evaluate"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may continue by deploying the exported model to [TensorRT](https://developer.nvidia.com/tensorrt) for optimized inference."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
